from psycopg2 import connect
from psycopg2.extensions import connection


class DBHandler:
    """数据库连接处理，可以为不同的数据库扩展连接方法。"""

    def __init__(
        self, host: str, port: int, database: str, user: str, password: str
    ) -> None:
        """
        初始化，传入数据库的通用连接配置。
        :param host: 主机地址
        :param port: 端口
        :param database: 数据库
        :param user: 登录用户名
        :param password: 登录密码
        """
        self.host = host
        self.port = port
        self.database = database
        self.user = user
        self.password = password

    def connect_postgresql(self) -> connection:
        """
        连接PostgreSQL数据库。
        :return: connect
        """
        conn = connect(
            host=self.host,
            port=self.port,
            database=self.database,
            user=self.user,
            password=self.password,
        )
        conn.cursor()
        return conn
